#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2022/2/10 17:26
# @Author  : CHENWang
# @Site    :
# @File    : plot_tool.py
# @Software: PyCharm


"""
0.5.1预计新增内容:
维护一个画图方法,之后可能会做成抽象基类
主要是画DataStruct的k线图, DataStruct加指标的图,以及回测框架的回测结果的图

https://blog.51cto.com/u_15287009/3227267
"""

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.gridspec as gridspec
from matplotlib import pyplot as plt
import matplotlib.animation as animation
import seaborn as sns
import squarify
from typing import Union
from functools import wraps

try:
    from pyecharts import Kline
except:
    from pyecharts.charts import Kline

import math
from quant_researcher.quant.project_tool import celebrity
from quant_researcher.quant.project_tool.logger.my_logger import LOG

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def plot_hq(df, y_col_list, x_col, mark_zero=False, all_in_one=None, normalize=None):
    LOG.info('The X axis is: %s', x_col)
    LOG.info('The Y axises are: %s', y_col_list)
    fig = plt.figure(figsize=(16, 8))
    x_s = df[x_col]
    noa = len(x_s)
    LOG.info('Points on x: %s', noa)

    if not all_in_one:
        how_many_y = len(y_col_list)
        for ii in range(how_many_y):
            ax = fig.add_subplot(how_many_y, 1, ii + 1)
            a_col = y_col_list[ii]
            ax.set_xlabel(x_col)
            ax.set_ylabel(a_col)
            y_s = df[a_col]
            ax.plot(x_s, y_s, markersize=12.0, label=ii)
            ax.grid(True, color='black')
            for jj in ax.get_xticklabels():
                jj.set_rotation(-90)
    else:
        ax = fig.add_subplot(111)
        ax.set_xlabel(x_col)
        ax.set_ylabel('Y')
        for ii in y_col_list:
            if normalize:
                max_v = max(df[ii])
                min_v = min(df[ii])
                y_s = (df[ii] - min_v) / (max_v - min_v)
            else:
                y_s = df[ii]
            ax.plot(x_s, y_s, markersize=12.0, label=ii)

        if mark_zero:
            y_s = np.zeros(noa)
            ax.plot(x_s, y_s, markersize=12.0, label='ZERO')

        handles, labels = ax.get_legend_handles_labels()
        ax.legend(handles[::-1], labels[::-1], bbox_to_anchor=(1.005, 1), loc='upper left', borderaxespad=0.)

        for ii in ax.get_xticklabels():
            ii.set_rotation(-90)
    # plt.savefig(FIG_PATH)
    plt.grid(True)
    plt.show()


def draw_a_line_and_scatter_two_points(df, x_col, y_col, x_1, x_2, test_chg, chg_name=None, fig_name=False,
                                       show_fig=True, info=None):
    fig = plt.figure(figsize=(16, 9))
    ax = fig.add_subplot(111)
    x_s = df[x_col]
    y_s = df[y_col]
    ax.plot(x_s, y_s)

    y_1 = df[df[x_col] == x_1][y_col].values[0]
    y_2 = df[df[x_col] == x_2][y_col].values[0]

    ax.scatter([x_1, x_2], [y_1, y_2], color='red')

    ax.plot([x_1, x_1], [y_1, y_2], color='black', linewidth=1, linestyle="--")
    ax.plot([x_2, x_2], [y_1, y_2], color='black', linewidth=1, linestyle="--")
    ax.plot([x_1, x_2], [y_1, y_1], color='black', linewidth=1, linestyle="--")
    ax.plot([x_1, x_2], [y_2, y_2], color='black', linewidth=1, linestyle="--")

    chg = (y_2 - y_1) / y_1
    chg_name = 'Chane' if chg_name is None else chg_name
    title_str = 'Name: %s\n%s: %s\n↓\n%s: %s\n%s：%s' % (fig_name, x_1, y_1, x_2, y_2, chg_name, celebrity.to_percent_str(chg))
    if info is not None:
        title_str += '\nInfo: %s' % info
    LOG.info('title str: %s', title_str)
    LOG.info('chg: %s, test_chg: %s', chg, test_chg)
    if chg_name == u'最大回撤率':
        assert math.fabs(chg) == math.fabs(test_chg)
    else:
        assert chg == test_chg
    ax.set_title(title_str, fontsize=12, color='r')

    for ii in ax.get_xticklabels():
        ii.set_rotation(90)
    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.grid(True)
    if fig_name:
        LOG.info('This figure is saved: %s', fig_name)
        plt.savefig(fig_name)
    if show_fig:
        plt.show()


def scatter_this(x_list, y_list, **kwargs):
    """
    画一个散点图

    :param x_list: 横坐标
    :param y_list: list或者list of list，纵坐标
    :param kwargs:
        - x_name，str，X轴名字，默认：None
        - y_name，str，Y轴名字，默认：None
        - fig_name，str，图像名字，默认：None
        - max_e_point_x，，，默认：None
        - max_e_point_y，，，默认：None
        - min_var_point_y，，，默认：None
        - min_var_point_x，，，默认：None
        - multiple_scatter: 如果纵坐标有多个的话，默认：None
    :return:
    """
    fig_name = kwargs.pop('fig_name', None)
    x_name = kwargs.pop('x_name', None)
    y_name = kwargs.pop('y_name', None)
    max_e_point_x = kwargs.pop('max_e_point_x', None)
    max_e_point_y = kwargs.pop('max_e_point_y', None)
    min_var_point_x = kwargs.pop('min_var_point_x', None)
    min_var_point_y = kwargs.pop('min_var_point_y', None)
    multiple_scatter = kwargs.pop('multiple_scatter', None)

    risk_free = 0.015  # 无风险利率设定为1.5%
    how_many_times_bigger = 1
    # how_many_times_bigger = 100
    if multiple_scatter is None:
        how_many_plot = 1
        y_list = [y_list]
        x_list = [x_list]
        y_name = [y_name]
        x_name = [x_name]
    else:
        how_many_plot = len(y_list)
    fig = plt.figure(figsize=(16, 8))
    for ii in range(how_many_plot):
        ax = fig.add_subplot(how_many_plot, 1, 1 + ii)
        this_x = np.array(x_list[ii]) * how_many_times_bigger
        this_y = np.array(y_list[ii]) * how_many_times_bigger
        ax.scatter(this_x, this_y, c=(this_y - risk_free) / this_x, marker='o')
        if min_var_point_x is not None and min_var_point_y is not None:
            plt.plot([min_var_point_x], [min_var_point_y], 'r*', markersize=12.0)  # 红星：标记最小方差组合
        if max_e_point_y is not None and max_e_point_x is not None:
            plt.plot([max_e_point_x], [max_e_point_y], 'b*', markersize=12.0)  # 蓝星：标记最大sharpe组合
        ax.set_xlabel(x_name[ii])
        ax.set_ylabel(y_name[ii])
        ax.set_title(label='Sharpe Ratio(Is it? Maybe)')
    plt.grid(True)
    if fig_name is not None and isinstance(fig_name, str):
        plt.savefig(fig_name)
    plt.show()

    # plt.cla()  # Clear axis
    # plt.clf()  # Clear figure
    # plt.close()  # Close a figure window


def plot_datastruct(_quotation_base, code=None):
    import webbrowser
    if code is None:
        path_name = '.' + os.sep + 'QA_' + _quotation_base.type + \
                    '_codepackage_' + _quotation_base.if_fq + '.html'
        kline = Kline('CodePackage_' + _quotation_base.if_fq + '_' + _quotation_base.type,
                      width=1360, height=700, page_title='QUANTAXIS')

        data_splits = _quotation_base.splits()

        for i_ in range(len(data_splits)):
            data = []
            axis = []
            for dates, row in data_splits[i_].data.iterrows():
                open, high, low, close = row[1:5]
                datas = [open, close, low, high]
                axis.append(dates[0])
                data.append(datas)

            kline.add(_quotation_base.code[i_], axis, data, mark_point=[
                "max", "min"], is_datazoom_show=True, datazoom_orient='horizontal')
        kline.render(path_name)
        webbrowser.open(path_name)
        LOG.info('The Pic has been saved to your path: %s' % path_name)
    else:
        data = []
        axis = []
        for dates, row in _quotation_base.select_code(code).data.iterrows():
            open, high, low, close = row[1:5]
            datas = [open, close, low, high]
            axis.append(dates[0])
            data.append(datas)

        path_name = '.' + os.sep + 'QA_' + _quotation_base.type + \
                    '_' + code + '_' + _quotation_base.if_fq + '.html'
        kline = Kline(code + '__' + _quotation_base.if_fq + '__' + _quotation_base.type,
                      width=1360, height=700, page_title='QUANTAXIS')
        kline.add(code, axis, data, mark_point=[
            "max", "min"], is_datazoom_show=True, datazoom_orient='horizontal')
        kline.render(path_name)
        webbrowser.open(path_name)
        LOG.info('The Pic has been saved to your path: %s' % path_name)


def QA_plot_save_html(pic_handle, path, if_open_web):
    """
    explanation:
        将绘图结果保存至指定位置

    params:
        * pic_handle ->:
            meaning: 绘图
            type: null
            optional: [null]
        * path ->:
            meaning: 保存地址
            type: null
            optional: [null]
        * if_open_web ->:
            meaning: 是否调用浏览器打开
            type: bool
            optional: [null]

    return:
        None

    demonstrate:
        Not described

    output:
        Not described
    """

    pic_handle.render(path)
    if if_open_web:
        webbrowser.open(path)
    else:
        pass
    LOG.info('The Pic has been saved to your path: %s' % path)


class GridFigure:
    """
    使用网格视图
    """

    def __init__(self, rows, cols):
        self.rows = rows
        self.cols = cols
        self.fig = plt.figure(figsize=(14, rows * 7))
        self.gs = gridspec.GridSpec(rows, cols, wspace=0.4, hspace=0.3)
        self.curr_row = 0
        self.curr_col = 0

    def next_row(self):
        if self.curr_col != 0:
            self.curr_row += 1
            self.curr_col = 0
        subplt = plt.subplot(self.gs[self.curr_row, :])
        self.curr_row += 1

    def next_cell(self):
        if self.curr_col >= self.cols:
            self.curr_row += 1
            self.curr_col = 0
        subplt = plt.subplot(self.gs[self.curr_row, self.curr_col])
        self.curr_col += 1
        return subplt

    def close(self):
        plt.close(self.fig)
        self.fig = None
        self.gs = None


def customize(func):
    """
    修饰器，设置输出图像内容与风格
    """

    @wraps(func)
    def call_w_context(*args, **kwargs):
        set_context = kwargs.pop("set_context", True)
        if set_context:
            color_palette = sns.color_palette("colorblind")
            with plotting_context(), axes_style(), color_palette:
                sns.despine(left=True)
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)

    return call_w_context


def plotting_context(
        context: str = "notebook", font_scale: float = 1.5, rc: dict = None
):
    """
    创建默认画图板样式

    参数
    ---
    :param context: seaborn 样式
    :param font_scale: 设置字体大小
    :param rc: 配置标签
    """
    if rc is None:
        rc = {}

    rc_default = {"lines.linewidth": 1.5}

    # 如果没有默认设置，增加默认设置
    for name, val in rc_default.items():
        rc.setdefault(name, val)

    return sns.plotting_context(context=context, font_scale=font_scale, rc=rc)


def axes_style(style: str = "darkgrid", rc: dict = None):
    """
    创建默认轴域风格

    参数
    ---
    :param style: seaborn 样式
    :param rc: dict 配置标签
    """
    if rc is None:
        rc = {}

    rc_default = {}

    for name, val in rc_default.items():
        rc.set_default(name, val)

    return sns.axes_style(style=style, rc=rc)


def print_table(table: Union[pd.Series, pd.DataFrame], name: str = None, fmt: str = None):
    """
    设置输出的 pandas DataFrame 格式

    参数
    ---
    :param table: 输入表格
    :param name: 表格列名设置
    :param fmt: 设置表格元素展示格式，譬如设置 '{0:.2f}%'，那么 100 展示出来就是 '100.00%'
    """
    if isinstance(table, pd.Series):
        table = pd.DataFrame(table)

    if isinstance(table, pd.DataFrame):
        table.columns.name = name

    prev_option = pd.get_option("display.float_format")
    if fmt is not None:
        pd.set_option("display.float_format", lambda x: fmt.format(x))

    from IPython.display import display
    display(table)

    # from alphalens, seems useless
    # if fmt is not None:
    #     pd.set_option("display.float_format", prev_option)


def plot_ohlcv(data):
    # 成交量可视化
    # 绘制K线图+移动平均线+成交量

    data.index = pd.to_datetime(data.index)
    data.columns = data.columns.map(lambda x: x.capitalize())
    fig = plt.figure(figsize=(8, 6), dpi=100, facecolor="white")  # 创建fig对象
    gs = gridspec.GridSpec(2, 1, left=0.10, bottom=0.15, right=0.96, top=0.96, wspace=None, hspace=0, height_ratios=[3.5, 1])
    graph_KAV = fig.add_subplot(gs[0, :])
    graph_VOL = fig.add_subplot(gs[1, :])

    # 绘制K线图
    from mplfinance import original_flavor
    original_flavor.candlestick2_ochl(graph_KAV, data.Open, data.Close, data.High, data.Low,
                                      width=0.5, colorup='r', colordown='g')  # 绘制K线走势

    # 绘制移动平均线图
    data['Ma20'] = data.Close.rolling(window=20).mean()  # pd.rolling_mean(data.Close,window=20)
    data['Ma30'] = data.Close.rolling(window=30).mean()  # pd.rolling_mean(data.Close,window=30)
    data['Ma60'] = data.Close.rolling(window=60).mean()  # pd.rolling_mean(data.Close,window=60)

    graph_KAV.plot(np.arange(0, len(data.index)), data['Ma20'], 'black', label='M20', lw=1.0)
    graph_KAV.plot(np.arange(0, len(data.index)), data['Ma30'], 'green', label='M30', lw=1.0)
    graph_KAV.plot(np.arange(0, len(data.index)), data['Ma60'], 'blue', label='M60', lw=1.0)

    graph_KAV.legend(loc='best')
    graph_KAV.set_title(u"BTC-日K线")
    graph_KAV.set_ylabel(u"价格")
    graph_KAV.set_xlim(0, len(data.index))  # 设置一下x轴的范围
    graph_KAV.set_xticks(range(0, len(data.index), 100))  # X轴刻度设定 每15天标一个日期

    # 绘制成交量图
    graph_VOL.bar(np.arange(0, len(data.index)), data.Volume,
                  color=['g' if data.Open[x] > data.Close[x] else 'r' for x in range(0, len(data.index))])
    graph_VOL.set_ylabel(u"成交量")
    graph_VOL.set_xlabel("日期")
    graph_VOL.set_xlim(0, len(data.index))  # 设置一下x轴的范围
    graph_VOL.set_xticks(range(0, len(data.index), 100))  # X轴刻度设定 每15天标一个日期
    graph_VOL.set_xticklabels([data.index.strftime('%Y-%m-%d')[index] for index in graph_VOL.get_xticks()])  # 标签设置为日期

    # X-轴每个ticker标签都向右倾斜45度
    for label in graph_KAV.xaxis.get_ticklabels():
        label.set_visible(False)  # 隐藏标注 避免重叠

    for label in graph_VOL.xaxis.get_ticklabels():
        label.set_rotation(45)
        label.set_fontsize(10)  # 设置标签字体

    plt.show()
    plt.close()


def plot_treemap(df=None, title="", figsize=(8, 6), n_visible=5, **kwargs):
    """

    :param df: 每天都有一个treemap， 因此index为日期，columns为标签
                         A   B   C   D   E   F   G
            2020-01-01  10  20  25  35  10  25  45
            2020-01-02  25  20  25  35  10  25  30

    :param title: 画布标题
    :param figsize: 画布大小
    :param n_visible: treemap展示多少项
    :param kwargs:
        - int dpi: 默认为120
        - int duration: 默认为2
        - str filename: 默认为None, 支持输出格式.jpg(如果df只有一期数据), .gif(df包含多期)
    :return:
    """

    # 测试数据
    # 单期数据
    # df = pd.DataFrame([[10, 20, 25, 35, 10, 25, 45]], columns=["A", "B", "C", "D", "E", "F", "G"], index=['2020-01-01'])
    # 多期数据
    # df = pd.DataFrame([[10, 20, 25, 35, 10, 25, 45], [25, 20, 25, 35, 10, 25, 30]], columns=["A", "B", "C", "D", "E", "F", "G"], index=['2020-01-01', '2020-01-02'])

    # 创建画布
    dpi = kwargs.pop('dpi', 120)
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

    # 调整spines
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    ax.spines["bottom"].set_visible(False)

    # 设置颜色
    # 我们需要设置颜色，以保持所有内容的一致性。下图按顺序显示了所选的颜色
    color_map = ['#2E91E5', '#1CA71C', '#DA16FF', '#B68100', '#EB663B', '#00A08B', '#FC0080', '#6C7C32', '#862A16', '#620042', '#DA60CA', '#0D2A63'] * 100

    def plot_frame(i):
        ax.clear()
        date = df.index[i]
        dfdata = pd.DataFrame(df.loc[date, :])
        dfdata.columns = ["values"]
        dfdata["color"] = [color_map[i] for i in range(len(dfdata))]
        dfdata = dfdata.sort_values("values", ascending=False)

        dfused = dfdata.iloc[:n_visible, :]
        xdata = dfused.index.tolist()
        ydata = dfused["values"].tolist()
        colors = dfused["color"].tolist()

        squarify.plot(sizes=ydata,  # 指定绘图数据
                      label=xdata,  # 指定标签
                      color=colors,  # 指定自定义颜色
                      alpha=0.5,  # 指定透明度
                      value=ydata,  # 添加数值标签
                      edgecolor='white',  # 设置边界框为白色
                      linewidth=3,  # 设置边框宽度为3
                      ax=ax,
                      text_kwargs={"size": 8, "color": "black"},
                      zorder=0
                      )

        # 辅助设置
        ax.set_title(title, color="black", fontsize=12)
        ax.text(0.12, 0.92, str(date), va="center", ha="center", alpha=1.0, color="white", size=12, transform=ax.transAxes)
        plt.axis('off')
        plt.tick_params(top='off', right='off')

        # # 使用 Plotly Express 创建树状图， Plotly 构建的树状图是交互式的
        # import plotly.express as px
        # fig = px.treemap(df, path=['labels'], values='values', width=800, height=400)
        # fig.update_layout(
        #     treemapcolorway=colors,  # defines the colors in the treemap
        #     margin=dict(t=50, l=25, r=25, b=25))
        # fig.show()

        return plt

    filename = kwargs.pop('filename', None)
    if len(df.index) == 1:  # 只有一期不用做动画直接画图
        plt_res = plot_frame(0)
        if filename is not None and filename.endswith(".jpg"):
            plt_res.savefig(filename, dpi=300, bbox_inches='tight')
            plt_res.close()

    else:  # 多期则需要做动画
        duration = kwargs.pop('duration', 2)
        my_animation = animation.FuncAnimation(fig, plot_frame, frames=range(0, len(df)), interval=int(duration * 1000))

        if filename.endswith('.gif'):
            my_animation.save(filename, writer='imagemagick')
        else:
            raise NotImplementedError

    return


def plot_active_bar(df=None):
    # https://blog.csdn.net/XRTONY/article/details/114428048

    # 读取和清洗测试数据
    gdp = pd.read_excel("API_NY.GDP.MKTP.CD_DS2_zh_excel_v2_4671150.xls", header=3)  # https://api.worldbank.org/v2/zh/indicator/NY.GDP.MKTP.CD?downloadformat=excel
    # 筛选：去掉世界、一些地区性的数据
    def is_country(x, fields):
        filter_list = [True for field in fields if field in str(x)]
        if sum(filter_list) >= 1:
            return False
        else:
            return True
    fields = ["世界", "收入国家", "地区", "南亚", "组织成员", "人口", "北美", "联盟", "IBRD", "IDA", "重债穷国", 'nan', '的情況下']
    gdp["is_country"] = gdp.apply(lambda x: is_country(x["Country Name"], fields), axis=1)
    df = gdp[gdp["is_country"] == True]

    datas = []
    for year in range(1960, 2021):
        year = str(year)
        df.sort_values(year, inplace=True, ascending=False)
        data = df[0:15]  # 排序，取前15名
        data.sort_values(year, inplace=True, ascending=True)
        data[year] = data[year] / (10 ** 11)
        datas.append([year, data[year].tolist(), data["Country Name"].tolist()])

    # 绘制动态图

    class Plot(object):
        """docstring for Plot"""

        def __init__(self, data):
            fig, ax = plt.subplots(figsize=(12, 6))
            self.fig = fig
            self.ax = ax
            self.data = data

        def showGif(self, save_path, writer='imagemagick'):
            plt.cla()
            ani = animation.FuncAnimation(fig=self.fig,
                                          func=self.update,
                                          frames=len(self.data),
                                          init_func=self.init,
                                          interval=0.5,
                                          blit=False,
                                          repeat=False)
            # 不用imagemagick时，可以保存为html
            ani.save(save_path, writer=writer, fps=3)  #

        def init(self):
            bar = self.ax.barh([], [], color='#0CD9F1')
            return bar

        def update(self, i):
            self.ax.cla()
            data = self.data[i]
            x = data[1]
            y = data[2]
            year = data[0]

            bars = []
            for k in range(len(x)):
                tmp = y[k]
                if "v2" in sys.version:
                    tmp = y[k].encode("utf-8")
                if tmp in ["中国"]:
                    bar = self.ax.barh(k, x[k], color='r')
                else:
                    bar = self.ax.barh(k, x[k], color='#0CD9F1')
                bars.append(bar)
            # 添加数据标签
            for rect in bars:
                rect = rect[0]
                w = rect.get_width()
                self.ax.text(w, rect.get_y() + rect.get_height() / 2, '%.1f' % float(w), ha='left', va='center')

            # 设置Y轴刻度线标签
            self.ax.set_title(year)
            self.ax.set_yticks(range(len(y)))
            self.ax.set_yticklabels(y)
            if "v2" in sys.version:
                self.ax.set_xlabel("GDP(百亿)".decode("utf-8"))
            else:
                self.ax.set_xlabel("GDP(百亿)")

    plot = Plot(datas)
    if "v2" in sys.version:  # 如果你是用python2运行
        plot.showGif("gdp.gif", writer="imagemagick")

    if "v3" in sys.version:  # r如果你是用python3运行
        plot.showGif("gdp.html", writer="html")

    return


if __name__ == '__main__':
    # 测试plot_treemap
    # jpg_file = "treemap_dance.jpg"
    # gif_file = "treemap_dance.gif"
    # plot_treemap(filename=gif_file)

    # plot_active_bar
    plot_active_bar()
